# -*- coding: utf-8 -*-
"""
@Env 
@Time 2024/9/3 上午11:07
@Author yzpang
@Function: 
"""
import os
from multiprocessing import Queue, Process, set_start_method

import torch

from ..entity.request import ModelMessage
from ..entity.base import ProcessInfo
from ..config.registry import model_registry
from modelserver.configs.model_config import MODEL_ROOT_PATH
from modelserver.configs.base_config import get_logger

logger = get_logger()


def list_models():
    for key in list(model_registry.keys()):
        if not check_model_worker(key):
            del model_registry[key]

    return list(model_registry.keys())


def load_model(model: ModelMessage):
    model_name = model.modelName
    if model_name in model_registry:
        raise KeyError(f"模型{model_name}已存在")
    model_path = os.path.join(MODEL_ROOT_PATH, model.modelPath)
    if 'auto' == model.device:
        model.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    kwargs = {"problem_type": model.classificationType, 'device': model.device}
    # 启动新进程加载模型
    inbox = Queue()
    outbox = Queue()

    process = Process(target=run_model_worker,
                      name="model_worker_" + model_name,
                      args=(model_name, model_path, model.modelFrame, inbox, outbox),
                      kwargs=kwargs)
    process.start()
    model_registry[model_name] = ProcessInfo(process=process, inbox=inbox, outbox=outbox)
    return model_name


def unload_model(model: ModelMessage):
    """卸载模型"""
    model_name = model.modelName
    # 首先判断模型是否在运行
    if model_name not in model_registry:
        raise KeyError(f"模型{model_name}不存在")
    process_info = model_registry[model_name]
    # 向模型进程发送停止信号
    # todo 队列阻塞,使用os.kill
    process = process_info.process
    process.terminate()
    del model_registry[model_name]
    return model_name


def run_model_worker(model_name: str,
                     model_path: str = "",
                     model_frame: str = "",
                     inbox: Queue = None,
                     outbox: Queue = None,
                     **kwargs):
    """
    创建新进程启动模型
    :param model_name: 模型名称
    :param model_path: 模型路径
    :param model_frame: 模型加载框架(paddle, transformers)
    :param inbox: 接收主进程数据的队列
    :param outbox: 发送数据到主进程的队列
    :param kwargs: 备用
    """
    from modelserver import model
    work_class = getattr(model, model_frame)
    device = kwargs.get('device')
    model_worker = work_class(model_name, model_path, device)
    model_worker.load_model()

    while True:
        message = inbox.get()
        if message == 'STOP':
            break
        else:
            result = model_worker.handle_message(message)
            outbox.put(result)
        logger.info(message)


def check_model_worker(model_name: str):
    """ 检查模型进程是否有效 """
    process_info = model_registry[model_name]
    process = process_info.process
    logger.info(f"{model_name} status: {process.is_alive()}")
    return process.is_alive()

